#include "NScriptTraverser.h"
#include "NScriptLoader.h"
#include "nsNodes.h"
#include "neiderra/utils/CommonFunctions.h"
#include "neiderra/utils/Logger.h"
#include "neiderra/utils/vm/VirtualMachine.h"
#include "neiderra/core/io/IOStream.h"

#include <cstdlib>
#include <string>
#include <vector>
#include <stack>
#include <algorithm>

using namespace neiderra::utils;
using namespace neiderra::utils::script;
using namespace std;

#define log(s) core->getLogger()->log(s);
#define trace(s) core->getLogger()->traceBlock(s)
#define leave core->getLogger()->leaveBlock()

iNScriptTraverser::iNScriptTraverser(Core* core, Symtable* sym)
: Base(core), iStreamable()
{

}

iNScriptTraverser::~iNScriptTraverser()
{
}

nsAsmGenerator::nsAsmGenerator(Core* core, Symtable* sym)
: iNScriptTraverser(core, sym)
{
	_symtable = sym;

	reset();
}

void nsAsmGenerator::traverseFrom(nsNode* node)
{
	trace("nsAsmGenerator::traverseFrom");

	reset();
	node->traverse(this);

	leave;
}

void nsAsmGenerator::reset()
{
	reg_offset = 0;
	utils::memset(_regs, 0, 16);

	enterScope(SCOPE_FREE);

	setCurrentClass(0);
	setCurrentFunction(0);
}

nsAsmGenerator::~nsAsmGenerator()
{
}

void nsAsmGenerator::traverse(nsNode* n, TraverseType t)
{
	trace("nsNode");
	leave;
}

void nsAsmGenerator::traverse(nsNodes* nodes, TraverseType t)
{
	trace("nsNodes");

	const unsigned count = nodes->count();

	log("count: ");
	log((int)count);

	for(unsigned i=0; i<count; i++)
	{
		nsNode* node = nodes->getByNum(i);
		if(node)
			node->traverse(this);
	}

	leave;
}

void nsAsmGenerator::traverseParams(nsNodes* params)
{
	trace("traverseParams");
	// + find free registers
	// + pop the params into them respectively (back-order)
	// + free registers
	// + db args
	// + st args, regs
	
	unsigned paramsNum = params->count();
	vector<unsigned> regsIndices;	// register numbers
	string code;
	
/*		// if this is a method, then
	if(isInScope(SCOPE_CLASS))
	{
		// fixme : no regs
		unsigned free_reg = find_free_reg();
		string rn = itoa(free_reg);
		// fixme: do the same everyware (open_reg, close_reg, ...)
		open_reg(free_reg);
		writeln("pop r" + rn, "get $this");
		writeln("db $this, " + itoa(VM_POINTER_SIZE));
		writeln("st $this, r" + rn + ", " + itoa(VM_POINTER_SIZE));
		close_reg(free_reg);
	}
*/

	int rnThis=0;	// register for this
	if(isInScope(SCOPE_CLASS))
	{
		int rnThis = find_free_reg();
		open_reg(rnThis);
		// fixme
		regsIndices.push_back(rnThis);
	}

	for(int i=0;i!=paramsNum;i++)
	{
		// find register for param
		int regNum = find_free_reg();
		open_reg(regNum);
		if(regNum < 0)
		{
			core->getLogger()->fatal("no free registers!");
			leave;
			throw ENoFreeRegisters();
		}
		regsIndices.push_back(regNum);
	}

	// if in method, pop this
	if(isInScope(SCOPE_CLASS))
	{
		writeln("pop r" + intToStr(rnThis) + ", " + intToStr(VM_POINTER_SIZE), "get $this");
	}

	// fixme: reg numbers

	// pop params
	for(int i=paramsNum-1;i>=0;i--)	
	{
		// get param size
		nsArg* param = (nsArg*)params->getByNum(i);
		unsigned argSize = param->getSize();
		
		// pop param
		code = "pop r";
		code += intToStr(i);
		code += ", ";
		code += intToStr(argSize);
		writeln(code, "get args...");
	}

	if(isInScope(SCOPE_CLASS))
	{
		writeln("db $this, " + intToStr(VM_POINTER_SIZE), "allocate mem for $this");
		writeln("st $this, r" + intToStr(rnThis) + ", " + intToStr(VM_POINTER_SIZE), "save $this");
	}

	// init params
	for(int i=paramsNum-1;i>=0;i--)	
	{
		// get param size
		nsArg* param = (nsArg*)params->getByNum(i);
		unsigned argSize = param->getSize();
		
		// pop param
		code = "db ";
		code += param->getName() + ", ";
		code += intToStr(argSize);
		writeln(code);
		
		code = "st ";
		code += param->getName() + ", r";
		code += intToStr(regsIndices.at(i));
		code += ", ";
		code += intToStr(argSize);
		writeln(code);
	}

	// free registers
	close_regs(regsIndices);
	writeln("");
	
	leave;
}

void nsAsmGenerator::traverse(nsFunc* nsf, TraverseType t )
{
	trace("nsFunc " + nsf->getName());
	setCurrentFunction(nsf);
	
	writeln("");
	writeln("proc " + nsf->getName() + ":", "proc begins here");
	
	if(nsf->getParams())
		traverseParams(nsf->getParams());

	if(nsf->getBody())	
		traverse(nsf->getBody());
	
	writeln("ret", nsf->getName() + " proc end");
	writeln("");

	setCurrentFunction(0);
	leave;
}

void nsAsmGenerator::traverse(nsClass* classnode, TraverseType t )
{
	trace("nsClass " + classnode->getName());
	
	enterScope(SCOPE_CLASS);
	
	traverse(&(classnode->methods));
	
	leaveScope();
	
	leave;
}

void nsAsmGenerator::traverse(nsBlock* block, TraverseType t )
{
	trace("nsBlock");

	traverse(block->stmts);

	leave;
}

void nsAsmGenerator::traverse(nsVarDec* var, TraverseType t )
{
	trace("nsVarDec " + var->getName());
	
	string code = "db ";
	code += var->getName() + ", ";

	if(var->isPointer() || var->isField())
	{
		code += intToStr(VM_POINTER_SIZE);
	}
	else
	{
		code += intToStr(var->getSize());
	}
	writeln(code, "var");
	writeln("");
	
	leave;
}

void nsAsmGenerator::traverse(nsVar* var, TraverseType t )
{
	trace("nsVar " + var->getName());
	
	// fixme ? : if(var->isField())
	if(isInScope(SCOPE_CLASS) && !var->isLocal())
	{
		string sz = intToStr(VM_POINTER_SIZE);
		//writeln("push $this, " + sz, "calculating " + var->getName() + " heap address");
		writeln("push " + var->getName() + ", " + sz);
		writeln("add", var->getName() + " address on the top");
        
		int regnum = find_free_reg();
		// fixme: if reg < 0

		string rn = utils::intToStr(regnum);
		sz = utils::intToStr(var->getSize());

		if(!isInScope(SCOPE_LVALUE))
		{
			writeln("pop r" + rn + ", " + sz, "save " + var->getName() + " address to r" + rn);
			writeln("pushm r" + rn + ", " + sz, "push " + var->getName() + " value to the top of the stack" );
		}

		
	}
	else if(!isInScope(SCOPE_LVALUE))
	{
		string code = "push ";
		code+=var->getName() + ", ";
		/*if(var->getType())
			code+=itoa(var->getType()->size());
		else
			code+=itoa(var->size());*/
		if(var->getType()->isEmbedded())
		{
			code += intToStr( var->getSize() );
		}
		else
			code += intToStr(VM_POINTER_SIZE);
		writeln(code);
	}

	leave;
}

void nsAsmGenerator::traverse(nsRetStmt* rets, TraverseType t )
{
	trace("nsRetStmt");
	
	enterScope(SCOPE_RETSTMT);

	if(rets->right)
		traverse(rets->right);

	writeln("");

	leaveScope();

	leave;
}

void nsAsmGenerator::traverse(nsExpr* exp, TraverseType traverseType )
{
	trace("nsExpr");

	if(traverseType == TT_NORMAL){
		log("TT_NORMAL");
	
		exp->getRight()->traverse(this);

		if(exp->getOperationType() == OT_ASSIGN)
		{
			enterScope(SCOPE_LVALUE);
			exp->getLeft()->traverse(this);

			leaveScope();
		}
		else
			exp->getLeft()->traverse(this);
		
	/*	int lsize = exp->left->type->size();
		int rsize = exp->right->type->size();*/

		// fixme : get real size
		int lsize = sizeof(int);
		int rsize = sizeof(int);
		
		int size=0;

		string code, comment;

		if(lsize != rsize)
			comment = "non-explisit type convertion";
		
		switch(exp->operation)
		{
			case OT_ASSIGN: 

				if(getScope() == SCOPE_CLASS)
				{
					int regnum = find_free_reg();
					string rn = utils::intToStr(regnum);
					string sz = utils::intToStr(lsize);

					writeln("popm r" + rn + ", " + sz, "save to mem");
					//writeln("pushm r" + rn + ", " + sz, "get value from mem");

				}
				else
				{
					if(exp->getLeft()->getType()->isEmbedded())
						// for local variables
					{
						code = "pop "; 
						string lvaluename = exp->getLeft()->getName();
						code += lvaluename + ", ";
						size = lsize > rsize ? lsize : rsize;
						code += intToStr(size);

						writeln(code, "put result into " + lvaluename);
					}
					else {			// for references
						int regnum = find_free_reg();
						string rn = utils::intToStr(regnum);
						string sz = utils::intToStr(lsize);
						
						writeln("pop r" + rn + ", " + sz);
						writeln("popm r" + rn + ", " + sz); 
					}
				}
				size = lsize;

				break;
			case OT_MULT: code="mul "; break;
			case OT_DIV: code="div "; break;
			case OT_PLUS: code="add "; break;
			case OT_MINUS: code="sub "; break;
		}
		
		// select bigger size in expr
		if(exp->getOperationType() != OT_ASSIGN){
			size = lsize > rsize ? lsize : rsize;
		
			code += intToStr(size);

			// todo : comments , ex: 'a*' or 'b+'
			writeln(code);
		}

		writeln("");
	}
	else if(traverseType == TT_GETVALUE)
	{
		log("TT_GETVALUE");
		leave;
		return;
	}

	leave;
}

void nsAsmGenerator::traverse(nsArg* arg, TraverseType t )
{
	trace("nsArg");
	
/*	// db arg
	string code = "db ";
	code += arg->name + ", ";
	code += itoa(arg->type->size());
	writeln(code, "passing param");
	
	// pop arg
	code = "pop ";
	code += arg->name + ", ";
	code += itoa(arg->type->size());
	writeln(code);
	
	writeln("");*/
	
/*	char buf[20];
	itoa(reg_offset, buf, 10);
	
	string code = "pop r";
	code += buf; 
	itoa(arg->type->size(), buf, 10);
	
	code += ", ";
	code += buf;
	reg_offset++;
	
	writeln(code);*/
	
	leave;
}

void nsAsmGenerator::traverse(nsAArg* arg, TraverseType t )
{
	trace("nsAArg");
	
	traverse(arg->left);
	
/*	string code = "mov ";
	int regnum = find_free_reg();
	open_reg(regnum);*/
	
	leave;
}

void nsAsmGenerator::traverse(nsOperand* op, TraverseType t )
{
	trace("nsOperand");
	
	op->traverse(this);
/*	string code = "push ";
	code += op->name + ", ";
	code += itoa(op->size());
	writeln(code);*/
	
	leave;
}

void nsAsmGenerator::traverse(nsFCall* fc, TraverseType t )
{
	trace("nsFCall");
	
	traverse(fc->params);
	
	string code = "call ";
	code += fc->getName();
	
	writeln(code);

	leave;
}

void nsAsmGenerator::traverse(nsSelfRef* ref, TraverseType t )
{
	trace("nsSelfRef ");

	writeln("push $this, " + intToStr(VM_POINTER_SIZE));

	if(ref->getRef())
		ref->getRef()->traverse(this);

	leave;
}

void nsAsmGenerator::traverse(nsMethodCall* ref, TraverseType t )
{
	trace("nsRefMethod " + ref->getName());

	writeln("push " + ref->getName() + ", " + intToStr(VM_POINTER_SIZE), "calculating $this for " + ref->getName());
	
	writeln("add");
	
	// fixme: if no regs?
	int regnum = find_free_reg();
	string rn = utils::intToStr(regnum);
	string sz = utils::intToStr(ref->getSize());

	// get method program address
	writeln("pop r" + rn + ", " + sz);

	// fixme: push params

	traverse(ref->getParams());
	
	//...

	// fixme: not here, POP $THIS

	if(isInScope(SCOPE_CLASS))
		writeln("push $this, " + intToStr(VM_POINTER_SIZE));
	else
	{
		writeln("push " + ref->getName(), "ref method, not in class");
	}
	writeln("call r" + rn);

	/*		if(!isInScope(SCOPE_LVALUE))
		{
			int regnum = find_free_reg();
			string rn = utils::itoa(regnum);

			// fixme: size = 0
			writeln("popm r" + rn + ", " + sz);
			writeln("push r" + rn + ", " + sz, "dereferencing " + ref->getName());
		}
//	}*/

	leave;
}

void nsAsmGenerator::traverse(nsRef* ref, TraverseType t )
{
	trace("nsRef " + ref->getName());

	// todo: dereference

	// fixme: DEREFERENCE

	string code;

	if(isInScope(SCOPE_CLASS))
	{
		code += "push " + ref->getName();
		writeln(code + ", " + intToStr(VM_POINTER_SIZE));
		writeln("add", ref->getName() + " address on the top");
	}
	else
	{
		code += "push " + ref->getName();
		writeln(code + ", " + intToStr(VM_POINTER_SIZE), ref->getName() + " address");
	}
	
	if(ref->getRef()){
		writeln("add");

		int regnum = find_free_reg();
		string rn = utils::intToStr(regnum);

		// fixme: size = 0
		string sz = utils::intToStr(ref->getRef()->getSize());
		writeln("popm r" + rn + ", " + sz);
		writeln("push r" + rn + ", " + sz, "dereferencing " + ref->getName());

		traverse(ref->getRef());

		if(!isInScope(SCOPE_LVALUE))
		{
			int regnum = find_free_reg();
			string rn = utils::intToStr(regnum);

			// fixme: size = 0
			string sz = utils::intToStr(ref->getRef()->getSize());
			writeln("popm r" + rn + ", " + sz);
			writeln("push r" + rn + ", " + sz, "dereferencing " + ref->getName());
		}
	}

	leave;
}

void nsAsmGenerator::traverse(nsNumber*, TraverseType t )
{
	trace("nsNumber");
	leave;
}

/*void nsAsmGenerator::traverse(nsField* field)
{
	trace("nsField " + field->getName());
	leave;
}

void nsAsmGenerator::traverse(nsMethod* method)
{
	trace("nsMethod" + method->getName());
	leave;
}*/

void nsAsmGenerator::writeln(const std::string& s)
{
	this->getOutputStream()->writeln(s);
}

void nsAsmGenerator::writeln(const std::string& cmd, const std::string& comment)
{
	this->getOutputStream()->write(cmd.c_str());
	if(comment!="")
	{
		getOutputStream()->write("\t\t; ");
		getOutputStream()->writeln(comment.c_str());
	}
}



int nsAsmGenerator::find_free_reg()
{
	for(int i=0;i<VM_REGISTERS_NUM;i++){
		if(!_regs[i])
		{
//			_regs[i] = true;
			return i;
		}
	}
	return -1;
}

bool nsAsmGenerator::open_reg(unsigned n)
{
	if(_regs[n] == false){
		_regs[n] = true; 
		return true;
	}
	else return false;
}

bool nsAsmGenerator::close_reg(unsigned n)
{
	if(_regs[n] == true){
		_regs[n] = false; 
		return true;
	}
	else return false;
}

bool nsAsmGenerator::open_regs(vector<unsigned>& indices)
{
	bool isOk = true;
	for(vector<unsigned>::iterator i = indices.begin();
		   i != indices.end();
		   i++)
	{
		unsigned n = *i;
		isOk = open_reg(n);
	}
	return isOk;
}

bool nsAsmGenerator::close_regs(vector<unsigned>& indices)
{
	bool isOk = true;
	for(vector<unsigned>::iterator i = indices.begin();
	   i != indices.end();
	   i++)
	{
		unsigned n = *i;
		isOk = close_reg(n);
	}
	return isOk;
}

int nsAsmGenerator::find_reg_for(const std::string& symbol)
{
	return -1;
}

void nsAsmGenerator::setCurrentClass(nsClass* c)
{
	_currentClass = c;
}

void nsAsmGenerator::setCurrentFunction(nsFunc* f)
{
	_currentFunction = f;
}

nsClass* nsAsmGenerator::getCurrentClass()
{
	return _currentClass;
}

nsFunc* nsAsmGenerator::getCurrentFunction()
{
	return _currentFunction;
}
